package org.luxor.commons.websocket.component.exporter;

import org.apache.tomcat.util.res.StringManager;
import org.apache.tomcat.websocket.server.WsServerContainer;
import org.luxor.commons.websocket.annotation.WebSocketServer;
import org.springframework.context.ApplicationContext;
import org.springframework.util.Assert;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

import javax.websocket.DeploymentException;
import javax.websocket.Encoder;
import javax.websocket.server.ServerContainer;
import javax.websocket.server.ServerEndpointConfig;
import java.util.Arrays;
import java.util.LinkedHashSet;
import java.util.Set;

/**
 * WebSocketServer自动注册
 *
 * @author Mr.yan  @date 2021/8/25
 */
public class WebSocketServerExporter extends ServerEndpointExporter {
    private static final StringManager sm = StringManager.getManager(WsServerContainer.class);

    @Override
    public WsServerContainer getServerContainer() {
        return (WsServerContainer) super.getServerContainer();
    }

    @Override
    protected void registerEndpoints() {
        super.registerEndpoints();
        Set<Class<?>> endpointClasses = new LinkedHashSet<>();
        ApplicationContext context = getApplicationContext();
        if (context != null) {
            String[] endpointBeanNames = context.getBeanNamesForAnnotation(WebSocketServer.class);
            for (String beanName : endpointBeanNames) {
                endpointClasses.add(context.getType(beanName));
            }
        }
        for (Class<?> endpointClass : endpointClasses) {
            registerServer(endpointClass);
        }
    }

    private void registerServer(Class<?> endpointClass) {
        ServerContainer serverContainer = super.getServerContainer();
        Assert.state(serverContainer != null,
                "No ServerContainer set. Most likely the server's own WebSocket ServletContainerInitializer " +
                        "has not run yet. Was the Spring ApplicationContext refreshed through a " +
                        "org.springframework.web.context.ContextLoaderListener, " +
                        "i.e. after the ServletContext has been fully initialized?");
        try {
            if (logger.isInfoEnabled()) {
                logger.info("Registering @WebSocketServer class: " + endpointClass);
            }

            WebSocketServer annotation = endpointClass.getAnnotation(WebSocketServer.class);
            if (annotation == null) {
                throw new DeploymentException(
                        sm.getString("serverContainer.missingAnnotation",
                                endpointClass.getName()));
            }
            String path = annotation.value();

            // Validate encoders
            validateEncoders(annotation.encoders());

            // ServerEndpointConfig
            ServerEndpointConfig sec;
            Class<? extends ServerEndpointConfig.Configurator> configuratorClazz =
                    annotation.configurator();
            ServerEndpointConfig.Configurator configurator = null;
            if (!configuratorClazz.equals(ServerEndpointConfig.Configurator.class)) {
                try {
                    configurator = annotation.configurator().getConstructor().newInstance();
                } catch (ReflectiveOperationException e) {
                    throw new DeploymentException(sm.getString(
                            "serverContainer.configuratorFail",
                            annotation.configurator().getName(),
                            endpointClass.getClass().getName()), e);
                }
            }
            sec = ServerEndpointConfig.Builder.create(endpointClass, path).
                    decoders(Arrays.asList(annotation.decoders())).
                    encoders(Arrays.asList(annotation.encoders())).
                    subprotocols(Arrays.asList(annotation.subprotocols())).
                    configurator(configurator).
                    build();

            serverContainer.addEndpoint(sec);
        } catch (DeploymentException ex) {
            throw new IllegalStateException("Failed to register @WebSocketServer class: " + endpointClass, ex);
        }
    }

    private static void validateEncoders(Class<? extends Encoder>[] encoders)
            throws DeploymentException {

        for (Class<? extends Encoder> encoder : encoders) {
            // Need to instantiate decoder to ensure it is valid and that
            // deployment can be failed if it is not
            @SuppressWarnings("unused")
            Encoder instance;
            try {
                encoder.getConstructor().newInstance();
            } catch (ReflectiveOperationException e) {
                throw new DeploymentException(sm.getString(
                        "serverContainer.encoderFail", encoder.getName()), e);
            }
        }
    }
}
